import sys, os
import pickle                       # 序列化导入
from navie_layer import *
from mnist import load_mnist
from PIL import Image
from collections import OrderedDict

sys.path.append(os.pardir)

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def softmax(x):
    max_value = np.max(x)
    exp_x = np.exp(x - max_value)
    sum_exp_x = np.sum(exp_x)
    y = exp_x / sum_exp_x

    return y

# 显示MNist数据集图像
def img_show(img):
    pil_img = Image.fromarray(np.uint8(img))
    pil_img.show()

# img_show(x_train[0].reshape(28, 28))

# 获得测试集
def get_data():
    (x_train, t_train), (x_test, t_test) = load_mnist(normalize = True, flatten = True, one_hot_label = False)
    return x_test, t_test

# 获得已经训练好的神经网络参数
def init_network():
    with open("sample_weight.pkl", "rb") as f:
        network = pickle.load(f)

    return network

def predict(network, x):
    W1, W2, W3 = network['W1'], network['W2'], network['W3']
    b1, b2, b3 = network['b1'], network['b2'], network['b3']
    a1 = np.dot(x, W1) + b1
    z1 = sigmoid(a1)
    a2 = np.dot(z1, W2) + b2
    z2 = sigmoid(a2)
    a3 = np.dot(z2, W3) + b3
    y = softmax(a3)

    return y

def mean_squared_error(y, t):
    return 0.5 * np.sum((y - t) ** 2)

def cross_entropy_error(y, t):
    if y.ndim == 1:
        y = y.reshape(1, y.size)
        t = t.reshape(1, t.size)

    # 监督数据是one-hot-vector的情况下，转换为正确解标签的索引
    if t.size == y.size:
        t = t.argmax(axis=1)

    batch_size = y.shape[0]
    # 标签形式的交叉熵误差：
    # return -np.sum(np.log(y[np.arange(batch_size), t] + 1e-7)) / batch_size
    return -np.sum(np.log(y[np.arange(batch_size), t] + 1e-7)) / batch_size         # 1e-7为了避免产生一个-inf值

def numerical_diff(f, t):           # 导数
    h = 1e-4                        # 值过小会产生舍入误差
    return (f(t + h) - f(t - h))/(2 * h)

def numerical_gradient(f, x):       # 全导数（偏导数）
    h = 1e-4
    grad = np.zeros_like(x)

    iter = np.nditer(x, flags = ['multi_index'], op_flags = ['readwrite'])
    while not iter.finished:
        idx = iter.multi_index
        grad[idx] = (f(x[idx] + h) - f(x[idx] - h))/(2 * h)
        iter.iternext()

    return grad

def gradient_descent(f, init_x, lr = 0.01, step_num = 100):
    x = init_x
    for i in range(step_num):
        grad = numerical_gradient(f, x)
        x -= lr * grad
    return x

class TwoLayerNet:
    def __init__(self, input_size, hidden_size, output_seize, weight_init_std = 0.01):
        self.params = {}
        self.params['W1'] = weight_init_std * np.random.randn(input_size, hidden_size)
        self.params['b1'] = np.zeros(hidden_size)

        self.params['W2'] = weight_init_std * np.random.randn(hidden_size, output_seize)
        self.params['b2'] = np.zeros(output_seize)

        # 生成层
        self.layers = OrderedDict()
        self.layers['Affine1'] = Affine(self.params['W1'], self.params['b1'])
        self.layers['Relu1'] = Relu()

        self.layers['Affine2'] = Affine(self.params['W2'], self.params['b2'])

        self.lastlayers = SoftmaxWithLoss()

    def predict(self, x):
        # W1, W2 = self.params['W1'], self.params['W2']
        # b1, b2 = self.params['b1'], self.params['b2']
        #
        # a1 = np.dot(x, W1) + b1
        # z1 = sigmoid(a1)
        # a2 = np.dot(z1, W2) + b2
        #
        # return softmax(a2)
        for layer in self.layers.values():
            x = layer.forward(x)

        return x

    def loss(self, x, t):    # x: 输入数据，t: 监督数据
        y = self.predict(x)

        # return cross_entropy_error(y, t)
        return self.lastlayers.forward(y, t)


    def accuracy(self, x, t):
        y = self.predict(x)
        y = np.argmax(y, axis = 1)
        t = np.argmax(t, axis = 1)

        accuracy = np.sum(y == t) / float(x.shape[0])
        return accuracy

    def numerical_gradient(self, x, t):
        loss_W = lambda W: self.loss(x,t)    # ?

        grads = {}

        # 类内调用自己的方法必须试用self，此处的numerical_gradient为外部方法
        grads['W1'] = numerical_gradient(loss_W, self.params['W1'])
        grads['b1'] = numerical_gradient(loss_W, self.params['b1'])

        grads['W2'] = numerical_gradient(loss_W, self.params['W2'])
        grads['b2'] = numerical_gradient(loss_W, self.params['b2'])

        return grads

    def gradient(self, x, t):
        # forward
        self.loss(x, t)

        # backward
        dout = 1
        dout = self.lastlayers.backward(dout)
        layers = list(self.layers.values())
        layers.reverse()
        for layer in layers:
            dout = layer.backward(dout)

        grad = {}
        grad['W1'], grad['b1'] = self.layers['Affine1'].dW, self.layers['Affine1'].db
        grad['W2'], grad['b2'] = self.layers['Affine2'].dW, self.layers['Affine2'].db

        return grad


if __name__ == '__main__':

    (x_train, t_train), (x_test, t_test) = load_mnist(normalize = True, one_hot_label = True)
    train_loss_list = []
    train_acc_list = []
    test_acc_list = []

    # 超参数
    iters_num = 10000
    train_size = x_train.shape[0]
    batch_size = 100             # min-batch size
    learning_rate = 0.4          # 学习率

    # 平均每次epoch(随机遍历的数据集大小等于总数量)的重复次数
    iter_per_epoch = max(train_size / batch_size, 1)

    network = TwoLayerNet(input_size = 784, hidden_size = 50, output_seize = 10)

    for i in range(iters_num):
        # 获取mini-batch
        # batch_mask = np.arange(train_size)
        batch_mask = np.random.choice(train_size, batch_size)
        x_batch = x_train[batch_mask]
        t_batch = t_train[batch_mask]

        # 计算梯度
        # grad = network.numerical_gradient(x_batch, t_batch)     # 基于数值微分的前向传播算法
        grad = network.gradient(x_batch, t_batch)                 # 通过误差的反向传播法

        # 更新参数
        for key in ('W1', 'b1', 'W1', 'b2'):
            network.params[key] -= learning_rate * grad[key]

        # 计算学习过程中的损失函数
        loss = network.loss(x_batch, t_batch)
        train_loss_list.append(loss)

        # 识别每个epoch的识别精度
        if i % iter_per_epoch == 0:
            train_acc = network.accuracy(x_train, t_train)
            test_acc = network.accuracy(x_test, t_test)
            train_acc_list.append(train_acc)
            test_acc_list.append(test_acc)
            print("train acc: " + str(train_acc) + ", test acc: " + str(test_acc))
